//=============================================================================
//
//   Copyright (C) 2020 All rights reserved.
//
//   Filename:		tensor_ops.rs
//
//   Author: Wang Zhecheng - wangzhecheng@yeah.net
//
//   Date: 2020-10-03 01:23
//
//   Last Editors: Wang Zhecheng - wangzhecheng@yeah.net
//
//   Last modified:	2020-11-06 18:27
//
//   Description:
//
//=============================================================================
use super::tensor::{Tensor, TensorType};
use crate::dimension::DimType;
use crate::Real;
use nalgebra::allocator::Allocator;
use nalgebra::DefaultAllocator;
use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};

///
/// add & sub
///
macro_rules! bin_ops_impl {
    ($Trait:ident,$method:ident;$TraitAssign:ident,$method_assign:ident) => {
        impl<D: DimType, P: TensorType> $Trait for Tensor<D, P>
        where
            DefaultAllocator: Allocator<Real, D::Dim, D::Dim>,
        {
            type Output = Self;
            #[inline]
            fn $method(self, rhs: Self) -> Self::Output {
                let mut res = Self::Output::default();
                res.set_value(&self.tensor().$method(rhs.tensor()).clone_owned());
                res
            }
        }
        impl<D: DimType, P: TensorType> $TraitAssign for Tensor<D, P>
        where
            DefaultAllocator: Allocator<Real, D::Dim, D::Dim>,
        {
            #[inline]
            fn $method_assign(&mut self, rhs: Self) {
                self.set_value(&self.tensor().$method(rhs.tensor()).clone_owned());
            }
        }
    };
}
bin_ops_impl!(Add, add; AddAssign, add_assign);
bin_ops_impl!(Sub, sub; SubAssign, sub_assign);

///
/// neg
///
impl<D: DimType, P: TensorType> Neg for Tensor<D, P>
where
    DefaultAllocator: Allocator<Real, D::Dim, D::Dim>,
{
    type Output = Tensor<D, P>;
    #[inline]
    fn neg(self) -> Self::Output {
        let mut res = self;
        res.value_mut().neg_mut();
        res
    }
}

///
/// mul & div
///
macro_rules! scalar_ops_impl {
    ($Trait:ident,$method:ident;$TraitAssign:ident,$method_assign:ident) => {
        impl<D: DimType, P: TensorType> $Trait<Real> for Tensor<D, P> {
            type Output = Self;
            #[inline]
            fn $method(self, rhs: Real) -> Self::Output {
                let mut res = self.clone();
                for left in res.value_mut().as_mut_slice().iter_mut() {
                    *left = left.$method(rhs);
                }
                res
            }
        }
        impl<D: DimType, P: TensorType> $TraitAssign<Real> for Tensor<D, P> {
            #[inline]
            fn $method_assign(&mut self, rhs: Real) {
                for j in 0..self.value_mut().ncols() {
                    for i in 0..self.value_mut().nrows() {
                        unsafe {
                            self.value_mut()
                                .get_unchecked_mut((i, j))
                                .$method_assign(rhs)
                        };
                    }
                }
            }
        }
        impl<D: DimType, P: TensorType> $Trait<Tensor<D, P>> for Real {
            type Output = Tensor<D, P>;
            #[inline]
            fn $method(self, rhs: Tensor<D, P>) -> Self::Output {
                let mut res = rhs.clone();
                for j in 0..res.value_mut().ncols() {
                    for i in 0..res.value_mut().nrows() {
                        res[(i, j)] = self.$method(res[(i, j)]);
                    }
                }
                res
            }
        }
    };
}
scalar_ops_impl!(Mul, mul; MulAssign, mul_assign);
scalar_ops_impl!(Div, div; DivAssign, div_assign);

///
/// test
///
#[test]
fn add_sub() {
    use crate::arith::tensor::Stress;
    use crate::dimension::PlaneStress;
    let mut t1 = Stress::<PlaneStress>::new(1f32, 3f32, 3f32, 2f32);
    let t2 = Stress::<PlaneStress>::new(2f32, 4f32, 4f32, 3f32);
    let t1addt2 = t1 + t2;
    assert_eq!(t1addt2[(0, 0)], 3f32);
    assert_eq!(t1addt2[(0, 1)], 7f32);
    assert_eq!(t1addt2[(1, 0)], 7f32);
    assert_eq!(t1addt2[(1, 1)], 5f32);
    t1 += t2;
    assert_eq!(t1[(0, 0)], 3f32);
    assert_eq!(t1[(0, 1)], 7f32);
    assert_eq!(t1[(1, 0)], 7f32);
    assert_eq!(t1[(1, 1)], 5f32);
    t1 -= t2;
    assert_eq!(t1[(0, 0)], 1f32);
    assert_eq!(t1[(0, 1)], 3f32);
    assert_eq!(t1[(1, 0)], 3f32);
    assert_eq!(t1[(1, 1)], 2f32);
    let t1subt2 = t1 - t2;
    assert_eq!(t1subt2[(0, 0)], -1f32);
    assert_eq!(t1subt2[(0, 1)], -1f32);
    assert_eq!(t1subt2[(1, 0)], -1f32);
    assert_eq!(t1subt2[(1, 1)], -1f32);
}
#[test]
fn mul_div() {
    use crate::arith::tensor::Stress;
    use crate::dimension::PlaneStress;
    let mut t1 = Stress::<PlaneStress>::new(1f32, 3f32, 3f32, 2f32);
    let twomult1 = 2 as Real * t1;
    assert_eq!(twomult1[(0, 0)], 2f32);
    assert_eq!(twomult1[(0, 1)], 6f32);
    assert_eq!(twomult1[(1, 0)], 6f32);
    assert_eq!(twomult1[(1, 1)], 4f32);
    let t1mul2 = t1 * 2 as Real;
    assert_eq!(t1mul2[(0, 0)], 2f32);
    assert_eq!(t1mul2[(0, 1)], 6f32);
    assert_eq!(t1mul2[(1, 0)], 6f32);
    assert_eq!(t1mul2[(1, 1)], 4f32);
    t1 *= 2 as Real;
    assert_eq!(t1[(0, 0)], 2f32);
    assert_eq!(t1[(0, 1)], 6f32);
    assert_eq!(t1[(1, 0)], 6f32);
    assert_eq!(t1[(1, 1)], 4f32);
    let t2 = t1 / 2 as Real;
    assert_eq!(t2[(0, 0)], 1f32);
    assert_eq!(t2[(0, 1)], 3f32);
    assert_eq!(t2[(1, 0)], 3f32);
    assert_eq!(t2[(1, 1)], 2f32);
    t1 /= 2 as Real;
    assert_eq!(t1[(0, 0)], 1f32);
    assert_eq!(t1[(0, 1)], 3f32);
    assert_eq!(t1[(1, 0)], 3f32);
    assert_eq!(t1[(1, 1)], 2f32);
    let t2 = 2 as Real / t1;
    assert_eq!(t2[(0, 0)], 2f32 / 1f32);
    assert_eq!(t2[(0, 1)], 2f32 / 3f32);
    assert_eq!(t2[(1, 0)], 2f32 / 3f32);
    assert_eq!(t2[(1, 1)], 2f32 / 2f32);
}
#[test]
fn neg() {
    use crate::arith::tensor::Stress;
    use crate::dimension::PlaneStress;
    let t1 = Stress::<PlaneStress>::new(1f32, 3f32, 3f32, 2f32);
    let expected1 = Stress::<PlaneStress>::new(-1f32, -3f32, -3f32, -2f32);
    assert_eq!(-t1, expected1);
}
